# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import logging
import sys
import matplotlib.pyplot as plt

import pandas

from coin.base.book.types import TradeSide
from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.base.param_util import to_list

from coin.exchange.base.order_gateway import OrderSide
from coin.exchange.okex_futures.kr_rest.currency import OkexCurrency
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.exchange.bitfinex_v2.kr_rest.product import BitfinexProduct
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.gdax.kr_rest.product import GdaxProduct

from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.calculator import MovingSum
import coin.strategy.mm.fastfeature.feed_manager as fmgr
from coin.strategy.mm.fastfeature.util import get_prefix


class FeatureContainer(object):
  def __init__(self, products):
    self._products = products
    self._ready_products = set()
    self._ready = False
    self._fields = {}
    self._exch_cumtrade = {}
    self._exch_cumabstrade = {}

  def fillup_prefix(self, prefix, fields, timestamp):
    fields = {(prefix + "_" + key): value for key, value in fields.items()}
    fields['timestamp'] = timestamp
    return fields

  def get_raw_fields(self, prefix, book):
    n = 10
    bnames = ["bid%d" % (n - i - 1) for i in range(n)] + \
        ["ask%d" % i for i in range(n)]
    bnamemap = {}
    for i, bname in enumerate(bnames):
      bnamemap[bname] = "%s_%s" % (i, bname)

    fields = {}
    for side, method in [("bid", book.iteritems_bid), ("ask", book.iteritems_ask)]:
      for i, (price, price_qty) in enumerate(method()):
        if i >= n:
          break
        iname = "%s%d" % (side, i)
        name = bnamemap[iname]
        fields[iname + '_price'] = price
        fields[iname + '_qty'] = price_qty.qty
    fields = self.fillup_prefix(prefix, fields, book.timestamp)
    return fields

  def get_raw_trade_fields(self, prefix, trade):
    if prefix not in self._exch_cumtrade:
      self._exch_cumtrade[prefix] = 0
      self._exch_cumabstrade[prefix] = 0

    assert trade.side in [TradeSide.TRADE_BUY_SIDE, TradeSide.TRADE_SELL_SIDE], trade.side
    signed_qty = trade.qty * (1 if trade.side == TradeSide.TRADE_BUY_SIDE else -1)

    self._exch_cumtrade[prefix] += signed_qty
    self._exch_cumabstrade[prefix] += trade.qty
    cum_trade_qty = self._exch_cumtrade[prefix]
    cum_abs_trade_qty = self._exch_cumabstrade[prefix]

    fields = {
        'trade_price': trade.price,
        'trade_qty': signed_qty,
        'cum_trade_qty': cum_trade_qty,
        'cum_abs_trade_qty': cum_abs_trade_qty
    }
    fields = self.fillup_prefix(prefix, fields, trade.timestamp)
    return fields

  def update_book(self, product, book):
    prefix = get_prefix(product)

    if product in self._products:
      self._fields.update(self.get_raw_fields(prefix, book))
      if product not in self._ready_products:
        self._ready_products.add(product)
        print(set(self._products) - set(self._ready_products))
      if len(self._ready_products) == len(self._products) and not self._ready:
        print("ready")
        self._ready = True
    else:
      raise ValueError(product, self._products)

  def update_trade(self, product, trade):
    prefix = get_prefix(product)
    self._fields.update(self.get_raw_trade_fields(prefix, trade))


class TimedSampler(object):
  def __init__(self, sampling_period=None):
    self._last_sampled = 0
    self._sampling_period = sampling_period or (1 * 1e9)

  def try_sample(self, timestamp):
    if self._last_sampled + self._sampling_period < timestamp:
      self._last_sampled = timestamp
      return True
    else:
      return False

  def update_book(self, book):
    return self.try_sample(book.timestamp)

  def update_trade(self, exchange, trade):
    return self.try_sample(trade.timestamp)


class OnBookSampler(object):
  def __init__(self, product):
    self._product = product

  def update_book(self, product, book):
    return product == self._product

  def update_trade(self, product, trade):
    return False


class ExploreBook(object):
  def __init__(self, machine, products, dump_dir='.', postfix=None, start_ts=None):
    self._machine = machine
    self._products = products
    self._tick_sizes = (0.5, 0.01)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._fields_out = []
    self._fields_view = []
    self._feature_container = FeatureContainer(self._products)

    okex_quarter = OkexFuturesProduct.FromStr("BTC-USD.QUARTER", start_ts)
    okex_this_week = OkexFuturesProduct.FromStr("BTC-USD.THIS_WEEK", start_ts)
    okex_next_week = OkexFuturesProduct.FromStr("BTC-USD.NEXT_WEEK", start_ts)
    # self._sampler = TimedSampler()
    self._watch_product = okex_quarter
    self._sampler = OnBookSampler(okex_quarter)
    self._time_set = set()
    self._agg_book = {}
    self._feed_manager = fmgr.FeedManager(self.on_book, self.on_trade)
    print(products)

  def on_book_reset(self, book_builder_name, book_builder):
    self._feed_manager.set_book_reset(self._products, book_builder_name, book_builder)

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return

    intervalstamp = int(book.timestamp / 600e9) * 600e9
    if intervalstamp not in self._time_set:
      self._time_set.add(intervalstamp)
      print(pandas.DatetimeIndex([self._start_ts]))
      print(pandas.DatetimeIndex([intervalstamp])[0])

    if hasattr(book, "get_aggregated_book"):
      if product.exchange == 'Bitflyer':
        aggb = book.get_aggregated_book(50)
      else:
        aggb = book.get_aggregated_book(0.5)
      self._feature_container.update_book(aggb)
      if aggb.ask0().price <= aggb.bid0().price:
        import pdb
        pdb.set_trace()
    else:
      self._feature_container.update_book(product, book)
    self.try_view(product)

  def on_trade(self, product, trade):
    exchange = product.exchange
    if trade.timestamp < self._start_ts:
      return
    self._feature_container.update_trade(product, trade)

  def plot_prod(self, plotdf, subprod, prod):
    xtime = pandas.DatetimeIndex(plotdf['timestamp'])
    sample_qty_map = {
        'OkexFutures_btc_quarter': 100,
        'OkexFutures_btc_next_week': 50,
        'OkexFutures_btc_this_week': 50,
        'Bitmex_XBTUSD': 100000,
        'Bitflyer_FX_BTC_JPY': 2,
    }

    sample_qty = sample_qty_map[subprod]

    for i in range(10):
      plt.plot(xtime,
               plotdf['%s_ask%s_price' % (prod, i)],
               'r-',
               xtime,
               plotdf['%s_bid%s_price' % (prod, i)],
               'k-',
               linewidth=0.1)
      for markersize, qty, marker in [(2.0, sample_qty * 5, 'k.'), (1.0, sample_qty, 'g.'),
                                      (0.5, sample_qty * 0.1, 'r.')]:
        idx = plotdf['%s_ask%s_qty' % (prod, i)] > qty
        xidx = plotdf.index[idx]
        xidx = xtime[xidx]
        if len(xidx) > 0:
          plt.plot(xidx,
                   plotdf.loc[idx, '%s_ask%s_price' % (prod, i)],
                   marker,
                   markersize=markersize)
        idx = plotdf['%s_bid%s_qty' % (prod, i)] > qty

        xidx = plotdf.index[idx]
        xidx = xtime[xidx]
        if len(xidx) > 0:
          plt.plot(xidx,
                   plotdf.loc[idx, '%s_bid%s_price' % (prod, i)],
                   marker,
                   markersize=markersize)
    # import pdb; pdb.set_trace()
    if '%s_trade_price' % prod in plotdf.columns:
      plt.plot(xtime, plotdf['%s_trade_price' % prod], 'b-', linewidth=0.5)
      # plt.grid(True)
      plt.twinx()
      plt.plot(xtime, plotdf['%s_cum_trade_qty' % prod], 'g.-', linewidth=0.5, markersize=0.5)
    plt.gcf().autofmt_xdate()
    plt.title(prod)

  def try_view(self, product):
    if product == self._watch_product:
      from pylab import rcParams
      rcParams['figure.figsize'] = 16, 7
      self._fields_view.append(self._feature_container._fields.copy())
      if (len(self._fields_view) % 200 != 0):
        return

      plotdf = pandas.DataFrame(self._fields_view[-800:])
      summ = plotdf[plotdf.columns[plotdf.columns.str.find("abs_trade_qty") >= 0]].max()
      print(summ[(summ.index.str.find("BTC") >= 0)])
      # plt.subplot(212); plt.plot(sorted(plotdf['timestamp'].diff()))
      subprod = self._watch_product.subscription_symbol
      prod = get_prefix(self._watch_product)
      plt.subplot(211)
      self.plot_prod(plotdf, subprod, prod)
      plt.subplot(212)
      self.plot_prod(plotdf, 'Bitmex_XBTUSD', 'Bitmex_BTC-USD.PERPETUAL')
      plt.waitforbuttonpress(timeout=-1)
      plt.clf()


def get_products(machine, current_datetime):
  raise ValueError(machine)
  return products


def launch(date):
  # machine='feed-01.eu-west-1.aws',
  # machine = 'feed-01.ap-northeast-1.aws'
  for machine in [
      'feed-01.eu-west-1.aws',  #'feed-01.ap-northeast-1.aws'
  ]:
    start_time = datetime.datetime(2018, 4, date, hour=0)
    hours = 24
    products = get_products(machine, start_time)

    strategy = ExploreBook(machine=machine,
                           products=products,
                           dump_dir='dump',
                           postfix='_%s_%dH' % (start_time.strftime('%Y%m%d'), hours),
                           start_ts=start_time)

    run_from_archive(products,
                     strategy.on_book_reset,
                     start_time,
                     start_time + datetime.timedelta(hours=hours),
                     machine=[machine],
                     worker_id='1')
  return 0


def main(argv):
  # date = 12
  date = 19
  launch(date)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  try:
    sys.exit(main(sys.argv))
  except KeyboardInterrupt:
    sys.exit(1)
